import string

## A class used for automatically generating code associated with parsing and storing
#  a user input option.
class inop:

  ## The name of the option, which will be the base of its C++ variable name.
  _name = ""

  ## The default value for the option.
  #
  #  This text will be used for the argument(s) of the option variable's C++ constructor.
  _default = ""

  ## A description of the option, which will be turned into a Doxygen comment.
  _description = ""

  ## The C++ type of the variable that holds the option.
  _cxx_type = ""

  ## Constructs an input option, storing the name, type, defualt value, and description.
  def __init__(self, name, cxx_type, default, description):
    self._name = name
    self._cxx_type = cxx_type
    self._default = default
    self._description = description

## The master list of the available user input options, their types, default values,
#  and their descriptions.
inp_opts = [
      
  inop("active_energy_only", "bool", "false",
       """
       whether to include nuclear and core energies when computing the local energy
       """
      ),

  inop("agp_source", "std::string", "pcps::option_value::AGP_Source_None",
       """
       where the initial agp weights should be gotten from \\sa pcps::option_value
       """
      ),

  inop("arena_block_size", "int", "10000000",
       """
       size, in bytes, of the memory blocks allocated by the arena memory manager
       """
      ),

  inop("arena_buffer_size", "int", "1000000",
       """
       size, in bytes, of the buffer allocated around blocks by the arena memory manager
       """
      ),

  inop("begin_true_projector", "long int", "-1",
       """
       sample at which to switch to the true projector (i.e. no fixed nodes) in projector MC
       """
      ),

  inop("bfgs_length", "int", "20",
       """
       number of gradient vectors retained in the L-BFGS method
       """
      ),

  inop("boundary", "std::string", "pcps::option_value::Boundary_Open",
       """
       the boundary conditions for the lattice \\sa pcps::option_value
       """
      ),

  inop("burn_in_length", "long int", "10000",
       """
       number of burn in steps to apply to the markov chain on EACH thread
       """
      ),

  inop("cg_thresh", "double", "1e-6",
       """
       threshold for solving linear equations with the conjugate gradient method
       """
      ),

  inop("cholesky_thresh", "double", "1.0e-10",
       """
       The accuracy threshold for the Cholesky decomposition of the two electron integrals.
       """
      ),

  inop("correlator_size", "int", "2",
       """
       size of correlators automatically generated based on the lattice
       """
      ),

  inop("correlator_type", "std::string", "pcps::option_value::Correlator_Line",
       """
       the type of correlator to build on the lattice (i.e. line, square, ...) \\sa pcps::option_value
       """
      ),

  inop("damping", "double", "0.0",
       """
       damping applied to the first few updates of projected optimizations (larger damping means slower change)
       """
      ),

  inop("diag_rdm_damping", "double", "0.0",
       """
       how much to damp the diagonal contributions to the rdm importance sampling function
       """
      ),

  inop("do_perturb", "bool", "false",
       """
       whether to use perturbative long range pair correlators in the projected method
       """
      ),

  inop("dump_freq", "int", "1",
       """
       how frequently to dump the current wavefunction to a binary file
       """
      ),

  inop("energy_tolerance", "double", "1.0e-7",
       """
       energy tolerance used by the projected optimization method
       """
      ),

  inop("fci_shift", "double", "0.0",
       """
       shift applied to the Hamiltonian's diagonal in the FCI method
       """
      ),

  inop("fix_classical", "bool", "false",
       """
       whether to fix the correlators generated by the "classical" correlator source
       """
      ),

  inop("global_seed", "int", "-1",
       """
       Overall random seed.  If < 0, the current time will be used instead.
       """
      ),

  inop("guiding_noise", "double", "1e-4",
       """
       Magnitude of the noise to use in the guiding function.
       """
      ),

  inop("hamiltonian_type", "std::string", "pcps::option_value::Ham_Hubbard",
       """
       the type of Hamiltonian to use \\sa pcps::option_value
       """
      ),

  inop("heisenberg_b", "double", "0.0",
       """
       external field strength in the Heisenberg model in units of J
       """
      ),

  inop("hf_thresh", "double", "1.0e-9",
       """
       convergence threshold for the energy in hartree fock calculations
       """
      ),

  inop("hubbard_t", "double", "1.0",
       """
       the hopping parameter in the Hubbard Model ( \\f$ H_1 = -t * \\sum_{\\langle i j \\rangle} a_i^+ a_j \\f$ )
       """
      ),

  inop("hubbard_u", "double", "0.0",
       """
       the onsite repulsion in the Hubbard Model ( \\f$ H_2 = U * \\sum_i n_{i \\uparrow} n_{i \\downarrow} \\f$ )
       """
      ),

  inop("init_config", "std::vector<int>", "",
       """
       used to force the Monte Carlo walkers to start in a specific configuration
       """
      ),

  inop("init_config_thresh", "double", "-1000.0",
       """
       threshold used for rejecting initial configurations with small weights
       \\sa pcps::good_init_config_value
       """
      ),

  inop("join_trotter_clusters", "bool", "false",
       """
       whether to join together the clusters corresponding to different trotter operators
       when computing the wavefunction update in time evolution
       """
      ),

  inop("lattice_length_x", "int", "1",
       """
       extent of the lattice along the x dimension
       """
      ),

  inop("lattice_length_y", "int", "1",
       """
       extent of the lattice along the y dimension
       """
      ),

  inop("lattice_length_z", "int", "1",
       """
       extent of the lattice along the z dimension
       """
      ),

  inop("lattice_type", "std::string", "pcps::option_value::Lattice_Cartesian",
       """
       the lattice type \\sa pcps::option_value
       """
      ),

  inop("max_ac_length", "int", "100",
       """
       the maximum lenth for which to compute autocorrelations
       """
      ),

  inop("max_cg_iter", "int", "1000",
       """
       maximum number of conjugate gradient iterations allowed
       """
      ),

  inop("max_grad_norm", "double", "1.0e100",
       """
       in steepest descent, gradients larger than this cutoff will be ignored
       """
      ),

  inop("max_hf_iter", "int", "0",
       """
       Maximum number of HF iterations used to generate initial orbitals.
       If > 0, HF will be used to optimize the initial orbitals after they have been
       read from their source.
       """
      ),

  inop("max_macro_loop", "int", "20",
       """
       maximum number iterations for various optimization methods
       """
      ),

  inop("max_nr_iter", "int", "0",
       """
       Maximum number of Newton-Raphson optimization steps to apply in the projected
       optimization method.  Once these are exhausted, the method will revert to the
       sweep algorithm for any remaining iterations.
       """
      ),

  inop("move_type", "std::string", "pcps::option_value::Move_Standard",
       """
       the type of move to use when evolving the markov chain \\sa pcps::option_value
       """
      ),

  inop("n_act_orb", "int", "0",
       """
       number of active orbitals
       """
      ),

  inop("n_cor_orb", "int", "0",
       """
       number of core orbitals
       """
      ),

  inop("nparticles", "int", "2",
       """
       number of particles
       """
      ),

  inop("nthreads", "int", "1",
       """
       number of threads to use (not all optimization methods support threading)
       """
      ),

  inop("optimization", "std::string", "pcps::option_value::Opt_None",
       """
       the optimization method or type of calculation to apply \\sa pcps::option_value
       """
      ),

  inop("optimize_agp_weights", "bool", "false",
       """
       whether agp weight variables should be optimized
       """
      ),

  inop("optimize_corrs", "bool", "false",
       """
       whether correlator variables should be optimized
       """
      ),

  inop("optimize_orb_coeffs", "bool", "false",
       """
       whether orbital coefficient variables should be optimized
       """
      ),

  inop("optimize_pairing_matrix", "bool", "false",
       """
       whether pairing matrix variables should be optimized
       """
      ),

  inop("orb_source", "std::string", "pcps::option_value::Orb_Source_None",
       """
       where the initial orbital coefficients should be gotten from \\sa pcps::option_value
       """
      ),

  inop("overlap_shift", "double", "0.001",
       """
       shift to apply to the diagonal of the overlap matrix in stochastic reconfiguration
       """
      ),

  inop("pairing_source", "std::string", "pcps::option_value::Pairing_Source_None",
       """
       where the initial pairing matrix should be gotten from \\sa pcps::option_value
       """
      ),

  inop("parallel_type", "std::string", "pcps::option_value::Parallel_Site",
       """
       the way work is distributed across processes in the projected method \\sa pcps::option_value
       """
      ),

  inop("particle_spin", "double", "0.5",
       """
       the spin of a single particle (e.g. 0.5 for fermions)
       """
      ),

  inop("particle_type", "std::string", "pcps::option_value::Part_Fermion",
       """
       the type of particle (e.g. fermion, spin, ...) \\sa pcps::option_value
       """
      ),

  inop("pmc_tau", "double", "0.001",
       """
       timestep used in fixed phase projector Monte Carlo
       """
      ),

  inop("pop_control_freq", "int", "10",
       """
       frequency at which birth/death and overall population control is applied
       """
      ),

  inop("pqpq_guiding", "bool", "false",
       """
       whether to weight the elements of the local rdm by the [pq,pq] integrals
       """
      ),

  inop("precondition_sr", "bool", "false",
       """
       whether to use a diagonal preconditioner for the overlap matrix in stochastic reconfiguration
       """
      ),

  inop("print_freq", "int", "1",
       """
       how frequently to print the status of the optimization
       """
      ),

  inop("project_sz", "bool", "true",
       """
       whether to project the total spin of the wavefunction to a specific value
       """
      ),

  inop("properties", "bool", "false",
       """
       whether to compute additional wavefunction properties, beyond just the energy
       """
      ),

  inop("properties_sample_length", "long int", "10000",
       """
       number of samples to use when evaluating wavefunction properties other than the energy
       """
      ),

  inop("random_correlators", "bool", "false",
       """
       whether to randomize the non-fixed, non-ti correlators after they are read from their sources
       """
      ),

  inop("ref_type", "std::string", "pcps::option_value::Ref_Uniform",
       """
       the reference wavefunction that will be operated on by the correlator operators \\sa pcps::option_value
       """
      ),

  inop("s_half_orth", "bool", "false",
       """
       whether to use the S^(-1/2) orthonormal orbital basis when reading the psi3 integrals
       """
      ),

  inop("sample_length", "long int", "10000",
       """
       total number of samples to be taken, this will be divided evenly between the processes/threads
       """
      ),

  inop("scalar_type", "std::string", "pcps::option_value::Scalar_Real",
       """
       type of numbers to use, either real or complex \\sa pcps::option_value
       """
      ),

  inop("slightly_random_orbs", "bool", "false",
       """
       whether to add a little noise to the orbital coefficients after a Hartree-Fock calculation
       """
      ),

  inop("spin_to_fermion_dump", "bool", "false",
       """
       whether to, at the end of a spin calculation, convert the wavefunction to fermion
       form and dump it to a file
       """
      ),

  inop("sr_disk_block_size", "int", "100",
       """
       size, in megabytes, of the blocks used for storing wavefunction derivatives in stochastic reconfiguration
       """
      ),

  inop("sr_use_disk", "bool", "false",
       """
       whether to store the wavefunction derivatives on disk during the stochastic reconfiguration proceedure
       """
      ),

  inop("step_size", "double", "0.1",
       """
       step size used in various optimization methods
       """
      ),

  inop("symmetric_pairing", "bool", "false",
       """
       whether to constrain the pairing matrix to be symmetric during optimization
       """
      ),

  inop("tei_screening", "double", "1e-9",
       """
       screening threshold for the two electron integrals
       """
      ),

  inop("trial_energy", "double", "0.0",
       """
       trial energy used in fixed phase projector Monte Carlo
       """
      ),

  inop("uhf_bias", "std::vector<double>", "0",
       """
       The alpha vs. beta bias to apply during unrestricted hartree fock calculations.
       In the hubbard model, this should be a single number representing the bias between
       up and down spin occupations between the A and B sublattices.  For other
       hamiltonians, this vector should contain individual biases for each lattice site.
       """
      ),

  inop("uhf_bias_decay", "double", "1.0",
       """
       multiplier that reduces the uhf bias in each iteration of the uhf scf procedure
       """
      ),

  inop("uhf_desegragation", "double", "0.1",
       """
       parameter used to re-mix alpha and beta orbitals after a uhf calculation for the
       pfaffian initial guess
       """
      ),

  inop("verbose_print", "bool", "false",
       """
       whether to print a lot or a little during execution
       """
      ),

  inop("walkers_per_process", "int", "100",
       """
       target number of random walkers on each process for fixed phase projector Monte Carlo
       """
      ),

  inop("workdir", "std::string", '"./"',
       """
       the directory in which to write any files (WARNING, this defaults to ./)
       """
      ),

  inop("z_spin", "double", "0.0",
       """
       if projecting to a specific total Sz, the value of Sz to project to
       """
      )

]

#  inop("corr_sources", "std::set<std::string>", "",
#       """
#       sources from which to build the initial correlators \\sa pcps::option_value
#       """
#      ),


## The master list of the options derived from the user input options.
der_opts = [

  inop("nsites", "int", "-1",
       """the number of sites in the lattice (for fermions,
       alpha and beta sites are counted separately)
       """
      ),

  inop("configs_per_site", "int", "-1",
       """
       the number of configurations a single site
       can take on (2 for fermions, 3 for spin = 1, etc.)
       """
      ),

  inop("n_agp_weights", "int", "-1",
       """
       the number of agp weights
       """
      ),

  inop("n_rest_orb_coeff", "int", "-1",
       """
       the number of restricted orbital coefficients
       """
      ),

  inop("n_pairing_variables", "int", "-1",
       """
       the number of variables in the pairing matrix
       """
      ),

  inop("n_elec_a", "int", "0",
       """
       the number of alpha electrons
       """
      ),

  inop("n_elec_b", "int", "0",
       """
       the number of beta electrons
       """
      ),

]

all_opts = inp_opts + der_opts

# write the cpp file
f = open("src/input_auto_gen.cpp", "w")
f.write(
"""
// ***** WARNING:  This is an automatically generated file.  Do not edit this file directly. *****
// *****           Instead, edit the input_auto_gen.py file.                                 *****

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \\file input_auto_gen.cpp
///
/// \\brief   file holding implementations of functions associated with reading user input
///
/// \\sa input_auto_gen.py
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <src/pcps.h>
#include <src/input.h>
#include <src/input_auto_gen.h>
#include <src/mpi_interface.h>
#include <src/archive.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \\brief   Constructs the object that holds user input data.
///
///   Sets data members to their default values and calls the initialization function.
///
/// \\param[in]     argc   The number of command line arguments passed to the program.
/// \\param[in]     argv   The command line arguments (the input file name is assumed to be the
///                       second argument).
///
///////////////////////////////////////////////////////////////////////////////////////////////////
"""
)
f.write("pcps::Input::Input()\n")
f.write("\n")
f.write("  // set data members' default values\n")
f.write("  :")
to_write = ""
for opt in all_opts:
  to_write += "\n    _%s(%s)," % (opt._name, opt._default)
# replace final comma with a carriage return 
to_write = to_write[:-1]
to_write += "\n" 
f.write(to_write)
del to_write
f.write("\n")
f.write("// initialize the object\n")
f.write("{}\n")
f.write("\n")
f.write(
"""
///////////////////////////////////////////////////////////////////////////////////////////////////
/// \\brief   Sets a member's value using the remainder of the input command after the member name
///          has been read.
///
///          This is the default template behavior.
///          There are more specialized versions of this function for bools, vectors, and sets.
///
/// \\param[in]     member_name  name of the member being set
/// \\param[in,out] arch         an archive holding the member's value
/// \\param[out]    member       the member whose value will be set
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class T> void pcps::Input::set_member(const std::string & member_name,
                                               pcps::Archive & arch,
                                               T & member) {

  // read and broadcast the member's value
  pcps::mpi::read_and_bcast(arch, member, (boost::format("failed to read Input member \\"%s\\"") % member_name).str());

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \\brief   Chooses which member's value to set based on the member's name and calls the
///          appropriate member-setting function.
///
/// \\param[in]     member_name  name of the member being set
/// \\param[in,out] arch         an archive holding the member's value
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::Input::set_member(const std::string & member_name, pcps::Archive & arch) {

"""
)
for i, opt in enumerate(inp_opts):
  if i == 0:
    f.write('  if      ')
  else:
    f.write('  else if ')
  f.write('(member_name == "%s") this->set_member(member_name, arch, _%s);\n' %(opt._name, opt._name))
f.write('  else\n    throw pcps::Exception("unknown input option name \\"%s\\"") % member_name;\n')
f.write('\n')
f.write('}\n')
f.close()

# write the header file
f = open("src/input_auto_gen.h", "w")
f.write(
"""
// ***** WARNING:  This is an automatically generated file.  Do not edit this file directly. *****
// *****           Instead, edit the input_auto_gen.py file.                                 *****

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \\file input_auto_gen.h
///
/// \\brief   file holding the declaration of the input option class
///
/// \\sa input_auto_gen.py
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_INPUT_AUTO_GEN_HEADER
#define PCPS_INPUT_AUTO_GEN_HEADER

#include <src/pcps.h>

namespace pcps {

  // pre-declaration of the archive base class
  class Archive;

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \\brief   A class to hold user input options.
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  class Input {

    public:

      Input();

      void set_member(const std::string & member_name, pcps::Archive & arch);

      void verify();

      void initialize_derived_data();

    private:

      Input(const Input &); ///< the copy constructor is disabled

      Input &operator=(const Input &); ///< the assignment operator is disabled

      template<class T> void set_member(const std::string & member_name, pcps::Archive & arch, T & member);

      //void initialize_from_file(const std::string & filename);

      //void process_input_command(const std::vector<char> & cmd, const int line_num);

      //void bcast_input_command(std::vector<char> & cmd, int line_num);

    // declare data members
    private:

"""
)
for opt in all_opts:
  comment = string.split(string.strip(opt._description), '\n')
  if len(comment) > 0:
    f.write("      /// \\brief  %s\n" % (string.strip(comment[0])))
  for c in comment[1:]:
    f.write("      ///         %s\n" % (string.strip(c)))
  f.write("      %s _%s;\n" % (opt._cxx_type, opt._name))
  f.write("\n")
f.write("    // data retrieval functions\n")
f.write("    public:\n")
f.write("\n")
for opt in all_opts:
  f.write("      const %s & %s() const { return _%s; }\n" %(opt._cxx_type, opt._name, opt._name))
f.write("\n")
f.write(
"""
  }; // end class Input

} // end namespace pcps

#endif
"""
)
f.close()

